import jieba.posseg as pos
import numpy as np

writer = open('./data/dmcnn_data.txt', 'w')
convert_role = {'O': 0, "P": 1, 'J': 2, 'T': 3, 'L': 4, "N": 5}
convert_event = {'statement': 1, "emergency": 2, 'perception': 3, 'stateChange': 4, 'operation': 5, "action": 6,
                 'movement': 7}


def add_one_recoder(words, marks, index, role_type, event_type):
    sent = ''
    for i, (word, mark) in enumerate(zip(words, marks)):
        if i == index:
            sent += word + '/B'
            sent += ', ' if i != len(marks) - 1 else '\n'
        else:
            sent += word + '/' + mark
            sent += ', ' if i != len(marks) - 1 else '\n'
    writer.write(sent)
    writer.write(sent)
    writer.write(str(event_type) + '-' + str(role_type) + '\n\n')


def deal_sent(sent):
    # tokens
    tokens = sent.split()
    event_type, event_index = -1, -1
    words, labels, marks = [], [], []
    for i in range(len(tokens)):
        marks.append('A')
    for i in range(len(tokens)):
        splited_token = tokens[i].split('/')
        words.append(splited_token[0])
        labels.append(splited_token[1])
        if splited_token[1] in convert_event.keys():
            event_type = convert_event[splited_token[1]]
            event_index = i
    if event_index == -1: return
    marks[event_index] = 'T'
    for i, label in enumerate(labels):
        if label in convert_event.keys():
            continue
        elif label in convert_role.keys():
            if convert_role[label] != 0:
                add_one_recoder(words, marks, i, convert_role[label], event_type)
                # todo 解决数据集的样本不均衡问题
            elif np.random.rand() > 0.8:
                add_one_recoder(words, marks, i, convert_role[label], event_type)


def read_data(file_path):
    with open(file_path) as f:
        lines = f.readlines()
        for line in lines:
            deal_sent(line)


read_data('../data/data.txt')
